package org.fujay.commons.langchain4j.rag.augmentor.injector;

import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.rag.content.Content;
import lombok.extern.slf4j.Slf4j;
import org.fujay.commons.langchain4j.rag.augmentor.EnhancedContentInjector;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static dev.langchain4j.internal.Utils.*;
import static java.util.stream.Collectors.joining;

/**
 * @author slm
 * @description 增强的内容注入器，提供更灵活地格式化和更多的提示模板选项
 */
@Slf4j
public class EnhancedInjector implements EnhancedContentInjector {

    public static final PromptTemplate DEFAULT_PROMPT_TEMPLATE = PromptTemplate.from(
            """
                    {{userMessage}}
                    
                    参考以下信息回答问题:
                    {{contents}}""");

    /**
     * 高级的提示模板，包含内容相关性和来源信息
     */
    public static final PromptTemplate ADVANCED_PROMPT_TEMPLATE = PromptTemplate.from(
            """
                    {{userMessage}}
                    
                    参考以下按相关性排序的信息回答问题:
                    {{contents}}
                    
                    请基于提供的参考资料回答，如无法从参考资料中找到答案，请明确说明。""");

    /**
     * 包含详细元数据的提示模板
     */
    public static final PromptTemplate DETAILED_PROMPT_TEMPLATE = PromptTemplate.from(
            """
                    {{userMessage}}
                    
                    请使用以下参考资料回答上述问题。参考资料已按相关性排序，每条资料包含来源和日期等元数据:
                    {{contents}}
                    
                    在回答中，请引用您使用的资料来源。如果参考资料不包含问题的答案，请说明"参考资料中没有相关信息"。""");

    private final PromptTemplate promptTemplate;
    private final List<String> metadataKeysToInclude;
    private final boolean includeContentNumbers;
    private final int maxContentLength;

    /**
     * 创建默认配置的增强注入器
     */
    public EnhancedInjector() {
        this(ADVANCED_PROMPT_TEMPLATE, null, true, 1500);
    }

    /**
     * 创建完整自定义配置的增强注入器
     *
     * @param promptTemplate        提示模板
     * @param metadataKeysToInclude 要包含的元数据键
     * @param includeContentNumbers 是否包含内容编号
     * @param maxContentLength      最大内容长度限制
     */
    public EnhancedInjector(
            PromptTemplate promptTemplate,
            List<String> metadataKeysToInclude,
            boolean includeContentNumbers,
            int maxContentLength) {
        this.promptTemplate = getOrDefault(promptTemplate, ADVANCED_PROMPT_TEMPLATE);
        this.metadataKeysToInclude = copyIfNotNull(metadataKeysToInclude);
        this.includeContentNumbers = includeContentNumbers;
        this.maxContentLength = maxContentLength;
    }

    /**
     * 创建Builder实例
     */
    public static EnhancedInjectorBuilder builder() {
        return new EnhancedInjectorBuilder();
    }

    /**
     * 将内容注入到聊天消息中
     */
    @Override
    public ChatMessage inject(List<Content> contents, ChatMessage chatMessage) {
        if (contents.isEmpty()) {
            return chatMessage;
        }

        log.debug("注入 {} 条内容到消息中", contents.size());

        Prompt prompt = createPrompt(chatMessage, contents);

        if (chatMessage instanceof UserMessage message && isNotNullOrBlank(message.name())) {
            return prompt.toUserMessage(message.name());
        } else if (chatMessage instanceof SystemMessage) {
            return prompt.toSystemMessage();
        }

        return prompt.toUserMessage();
    }

    /**
     * 创建包含内容的提示
     */
    protected Prompt createPrompt(ChatMessage chatMessage, List<Content> contents) {
        Map<String, Object> variables = new HashMap<>();

        if (chatMessage instanceof UserMessage userMessage) {
            variables.put("userMessage", userMessage.singleText());
        } else if (chatMessage instanceof SystemMessage systemMessage) {
            variables.put("userMessage", systemMessage.text());
        }
        // 限制内容数量，避免超过模型上下文窗口
        List<Content> limitedContents = limitContentsByLength(contents, maxContentLength);
        if (limitedContents.size() < contents.size()) {
            log.info("由于长度限制，内容数量从 {} 减少到 {}", contents.size(), limitedContents.size());
        }

        variables.put("contents", format(limitedContents));
        return promptTemplate.apply(variables);
    }

    /**
     * 根据最大长度限制内容
     */
    protected List<Content> limitContentsByLength(List<Content> contents, int maxLength) {
        if (maxLength <= 0) {
            return contents;
        }

        List<Content> result = new ArrayList<>();
        int totalLength = 0;

        for (Content content : contents) {
            String text = content.textSegment().text();
            if (totalLength + text.length() <= maxLength) {
                result.add(content);
                totalLength += text.length();
            } else {
                // 尝试截取部分内容
                int remainingSpace = maxLength - totalLength;
                if (remainingSpace > 100) { // 至少保留100个字符才有意义
                    String truncatedText = text.substring(0, remainingSpace) + "...";
                    TextSegment truncatedSegment = TextSegment.from(
                            truncatedText,
                            content.textSegment().metadata());
                    result.add(Content.from(truncatedSegment));
                }
                break;
            }
        }

        return result;
    }

    /**
     * 格式化内容列表
     */
    protected String format(List<Content> contents) {
        if (includeContentNumbers) {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < contents.size(); i++) {
                if (i > 0) {
                    sb.append("\n\n");
                }
                sb.append("【参考资料").append(i + 1).append("】\n");
                sb.append(format(contents.get(i), i + 1));
            }
            return sb.toString();
        } else {
            return contents.stream()
                    .map(content -> format(content, 0))
                    .collect(joining("\n\n"));
        }
    }

    /**
     * 格式化单个内容
     */
    protected String format(Content content, int index) {
        TextSegment segment = content.textSegment();

        StringBuilder sb = new StringBuilder();
        sb.append(segment.text());

        // 添加相关性得分 - 由于当前Content类没有score()方法，暂时注释掉
        // 如果需要显示相关性得分，可以考虑使用自定义Content类或通过元数据传递
        /*
         * if (showRelevanceScore && content.score() != null) {
         * sb.append("\n相关性: ").append(String.format("%.2f", content.score()));
         * }
         */

        // 添加元数据
        if (!isNullOrEmpty(metadataKeysToInclude)) {
            String metadataStr = formatMetadata(segment.metadata());
            if (!metadataStr.isEmpty()) {
                sb.append("\n").append(metadataStr);
            }
        }

        return sb.toString();
    }

    /**
     * 格式化元数据
     */
    protected String formatMetadata(Metadata metadata) {
        if (isNullOrEmpty(metadataKeysToInclude)) {
            return "";
        }

        List<String> metadataEntries = new ArrayList<>();
        for (String key : metadataKeysToInclude) {
            String value = metadata.getString(key);
            if (value != null) {
                metadataEntries.add(key + ": " + value);
            }
        }

        return String.join("\n", metadataEntries);
    }


    /**
     * 增强注入器构建器
     */
    public static class EnhancedInjectorBuilder {
        private PromptTemplate promptTemplate;
        private List<String> metadataKeysToInclude;
        private boolean includeContentNumbers = true;
        private int maxContentLength = 1500;

        EnhancedInjectorBuilder() {
        }

        public EnhancedInjectorBuilder promptTemplate(PromptTemplate promptTemplate) {
            this.promptTemplate = promptTemplate;
            return this;
        }

        public EnhancedInjectorBuilder metadataKeysToInclude(List<String> metadataKeysToInclude) {
            this.metadataKeysToInclude = metadataKeysToInclude;
            return this;
        }

        public EnhancedInjectorBuilder includeContentNumbers(boolean includeContentNumbers) {
            this.includeContentNumbers = includeContentNumbers;
            return this;
        }

        public EnhancedInjectorBuilder maxContentLength(int maxContentLength) {
            this.maxContentLength = maxContentLength;
            return this;
        }

        public EnhancedInjector build() {
            return new EnhancedInjector(
                    this.promptTemplate,
                    this.metadataKeysToInclude,
                    this.includeContentNumbers,
                    this.maxContentLength);
        }
    }
}
